#!/usr/bin/python

# Copyright (C) 2009 ApplianSys Ltd.
# 
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, 
# MA  02110-1301, USA.

"""
This module allows a programme to execute a shell command and wait until
either:
a) The command completes.
b) A timeout is exceeded.
"""

import os
import subprocess
import datetime 
import time
import signal

POLL_INTERVAL = 0.2     # poll the running command every 0.2 seconds to see if
                        # it has finished yet
DEFAULT_TIMEOUT = 3     # give commands three seconds to execute, by default

class TimeoutException(Exception):
    """
    Just an exception to raise when a command doesn't complete in good time.
    """

    pass


def command_timeout(command, timeout=DEFAULT_TIMEOUT, env=os.environ, **kwargs):
    """
    This calls a shell command and either returns its return code, output 
    and stderr 

    -or- 

    kills it if it doesn't exit within 'timeout' seconds and throws a
    TimeoutException.

    @param command: The command, as a list.  As per subprocess.Popen.
    @param timeout: Timeout in seconds.
    @param env: The environment in which to execute the command.
    @param **kwargs: Any other arguments which need to be passed to Popen.

    >>> command_timeout(["sleep", "1"], 2)
    (0, '', '')
    >>> command_timeout(["sleep", "10"], 2)
    Traceback (most recent call last):
    ...
    TimeoutException: Command '['sleep', '10']' timed-out after 2 seconds.
    >>> command_timeout(["echo", "hello world"], 2)
    (0, 'hello world\\n', '')
    >>> command_timeout(["ls", "/doesntexist"], 2)
    (2, '', 'ls: cannot access /doesntexist: No such file or directory\\n')

    """

    start = datetime.datetime.now()
    process = subprocess.Popen(command, 
                               stdout=subprocess.PIPE, 
                               stderr=subprocess.PIPE,
                               env=env,
                               **kwargs
                               )
    while process.poll() is None:
        time.sleep(POLL_INTERVAL)
        now = datetime.datetime.now()
        if (now - start).seconds > timeout:
            os.kill(process.pid, signal.SIGKILL)
            os.waitpid(-1, os.WNOHANG)
            raise TimeoutException("Command '%s' timed-out after %s seconds."
                                   % (command, timeout))
    return (process.poll(), process.stdout.read(), process.stderr.read())


def _test():
    import doctest
    failures, x = doctest.testmod()
    if failures > 0:
        sys.exit(1)


if __name__ == "__main__":
    _test()

 
